package ltd.trackan.aitextdialogue.controller;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.extern.slf4j.Slf4j;
import ltd.trackan.aitextdialogue.model.deo.OpenAIRequest;
import ltd.trackan.aitextdialogue.model.entity.OpenAI;
import ltd.trackan.aitextdialogue.model.enumeration.LargeScaleModel;
import ltd.trackan.aitextdialogue.utils.WebClientUtils;
import ltd.trackan.common.model.entity.RestBean;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import javax.validation.Valid;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author wxa
 */
@RestController
@Tag(name = "大模型多轮对话")
@RequestMapping("/deep/seek")
@Slf4j
public class DeepSeekController {
    
    
    // Session与DeepSeek会话历史
    private static final ConcurrentHashMap<String, OpenAI> SESSION_DEEP_SEEK = new ConcurrentHashMap<>();
    
    // Session与SseEmitter
    private static final ConcurrentHashMap<String, SseEmitter> SESSION_EMITTER = new ConcurrentHashMap<>();
    
    private void sendErrorAndComplete(SseEmitter emitter, String message) {
        try {
            emitter.send(SseEmitter.event()
                    .name("error")
                    .data("{\"finish_reason\": \"stop\",\"message\": \"" + message + "\"}"));
        } catch (Exception e) {
            log.error("发送错误信息失败: {}", e.getMessage(), e);
        } finally {
            emitter.complete();
        }
    }
    
    
    @Operation(summary = "流式调用 DeepSeek API")
    @PostMapping("/stream")
    private SseEmitter stream(@RequestBody @Valid OpenAIRequest openAIRequest) {
        String session = openAIRequest.getSession();
        String model = openAIRequest.getModel();
        String message = openAIRequest.getMessage();
        synchronized (session) {
            // SSE
            SseEmitter emitter = new SseEmitter(60000 * 30L);
            if (session.isEmpty()) {
                sendErrorAndComplete(emitter, "用户标识为空");
                return emitter;
            }
            // SESSION对应SSE
            SESSION_EMITTER.put(session, emitter);
            // 构造请求头
            Map<String, String> requestHeaders;
            try {
                requestHeaders = LargeScaleModel.getHeaders(model);
            } catch (Exception e) {
                sendErrorAndComplete(emitter, e.getMessage());
                return emitter;
            }
            // 判断模型并获取模型的地址
            String endpoint;
            try {
                endpoint = LargeScaleModel.getEndpoint(model);
            } catch (Exception e) {
                sendErrorAndComplete(emitter, e.getMessage());
                return emitter;
            }
            // 创建或获取请求session
            OpenAI openAISession = SESSION_DEEP_SEEK.get(session);
            if (openAISession == null) {
                OpenAI temp = new OpenAI();
                SESSION_DEEP_SEEK.put(session, temp);
                openAISession = temp;
            }
            openAISession.setModel(model);
            // 判断是否在会话中，只支持一个session，按业务一个用户可以由多个session，或一个用户只能有一个session
            if (openAISession.isStartSession()) {
                sendErrorAndComplete(emitter, "会话已开始");
                return emitter;
            }
            openAISession.setStartSession(true);
            // 创建当前对话环境
            RestBean<String> addSessionLog = openAISession.addSessionLog(new OpenAI.MessageResult().role("user").content(message));
            if (!addSessionLog.isSuccess()) {
                sendErrorAndComplete(emitter, "创建对话历史失败，原因：" + addSessionLog.getMsg());
                return emitter;
            }
            
            // SSE会话结束，最终都会调用的方法
            emitter.onCompletion(() -> {
                SESSION_DEEP_SEEK.get(session).setStartSession(false);
                log.info(session + "的SSE会话结束");
            });
            emitter.onTimeout(() -> log.warn(session + "的SSE连接超时"));
            emitter.onError(e -> log.error(session + "的SSE连接错误：{}", e.getMessage(), e));
            
            // 当前推理内容
            StringBuilder reasoningContent = new StringBuilder();
            // 当前回答内容
            StringBuilder content = new StringBuilder();
            log.info("开始请求session：" + session + "|模型：" + model + "|发送消息：" + message);
            
            // 请求体
            Map<String, Object> body = new HashMap<>();
            // 模型
            body.put("model", model);
            // 将对话历史添加到请求体
            body.put("messages", new ArrayList<>(openAISession.getSessionLog()));
            // 流式对话
            body.put("stream", true);
            
            // 发送请求
            WebClientUtils.sendStream(endpoint, requestHeaders, body, chunk -> {
                try {
                    // log.info(chunk);
                    // 客户端断开连接或被主动终止
                    if (SESSION_EMITTER.get(session) == null || !SESSION_DEEP_SEEK.get(session).isStartSession()) {
                        emitter.complete();
                        return;
                    }
                    // // 助手回答完成
                    // if ("[DONE]".equals(chunk)) {
                    //     RestBean<String> addSessionLog1 = SESSION_DEEP_SEEK.get(session).addSessionLog(new OpenAI.MessageResult().role("assistant").content(content.toString()));
                    //     if (!addSessionLog1.isSuccess()) {
                    //         log.error("添加助手回答失败，原因：" + addSessionLog1.getMsg());
                    //     }
                    //     log.info("思维模型：" + reasoningContent);
                    //     log.info("回答内容：" + content);
                    //     SESSION_DEEP_SEEK.get(session).setStartSession(false);
                    //     emitter.complete();
                    //     return;
                    // }
                    // chunk 是 JSON 字符串
                    // {"id":"","object":"","created":,"model":"","system_fingerprint":"",
                    // "choices":[{"index":0,"delta":{"content":"我是"},"logprobs":null,"finish_reason":null}]}
                    
                    JSONObject jsonObject = JSON.parseObject(chunk);
                    if (jsonObject.containsKey("choices")) {
                        // 内容数组
                        JSONArray choices = jsonObject.getJSONArray("choices");
                        for (int i = 0; i < choices.size(); i++) {
                            log.info(choices.getJSONObject(i).toString());
                            JSONObject choice = choices.getJSONObject(i);
                            JSONObject delta = choice.getJSONObject("delta");
                            if (delta.containsKey("reasoning_content")) {
                                reasoningContent.append(delta.getString("reasoning_content"));
                            }
                            if (delta.containsKey("content")) {
                                content.append(delta.getString("content"));
                            }
                            // 需要应对WebClient请求后404之类的可能出现的问题
                            if (choices.getJSONObject(i).get("finish_reason") != null) {
                                RestBean<String> addSessionLog1 = SESSION_DEEP_SEEK.get(session).addSessionLog(new OpenAI.MessageResult().role("assistant").content(content.toString()));
                                if (!addSessionLog1.isSuccess()) {
                                    log.error("添加助手回答失败，原因：" + addSessionLog1.getMsg());
                                }
                                log.info("思维模型：" + reasoningContent);
                                log.info("回答内容：" + content);
                                SESSION_DEEP_SEEK.get(session).setStartSession(false);
                                emitter.complete();
                            }
                            try {
                                emitter.send(SseEmitter.event().data(choice));
                            } catch (Exception e) {
                                log.warn(session + "客户端断开");
                                SESSION_EMITTER.remove(session);
                                emitter.complete();
                            }
                        }
                    }
                } catch (Exception e) {
                    log.error("出现chunk异常，开始结束SSE推送:{}", e.getMessage(), e);
                    SESSION_EMITTER.remove(session);
                    emitter.complete();
                    return;
                }
            });
            return emitter;
        }
    }
    
    @Operation(summary = "终止对话")
    @GetMapping("/cancel/{session}")
    public void cancelSession(@PathVariable String session) {
        synchronized (session) {
            log.info(session + "终止对话：");
            // SSE完成
            SseEmitter emitter = SESSION_EMITTER.remove(session);
            if (emitter != null) {
                emitter.complete();
            }
            // 终止对话
            OpenAI openAI = SESSION_DEEP_SEEK.get(session);
            if (openAI != null) {
                openAI.setStartSession(false);
            }
        }
    }
    
    @Operation(summary = "清空对话历史，开始新的一轮对话")
    @GetMapping("/clear-history/{session}")
    public void clearHistory(@PathVariable String session) {
        synchronized (session) {
            log.info(session + "清空对话历史：");
            // 清除SSE绑定
            SseEmitter emitter = SESSION_EMITTER.remove(session);
            if (emitter != null) {
                emitter.complete();
            }
            // 清除历史记录
            SESSION_DEEP_SEEK.remove(session);
        }
    }
}